#!/usr/bin/python
#This script plots phred quality values and G+C % of a contig
#Usage: python plotContigQuality.py contig.fasta contig.qual MINQUAL
#Author: Jimmy Saw
#Last updated: 12/25/2012

import sys
from Bio import SeqIO
import matplotlib.pyplot as plt
import pylab
import numpy as np
from Bio.SeqUtils import GC

def fmt(f):
    st = '{0:.4}'.format(f)
    return st

seqfile = open(sys.argv[1], "rU")
qualfile = open(sys.argv[2], "rU")
MINQUAL = int(sys.argv[3])

sr = SeqIO.read(seqfile, "fasta")
qr = SeqIO.read(qualfile, "qual")

tmp = sys.argv[1].split('.')
contigname = tmp[0]

avgqual = 0
qualrecs = 0
slist = []
qlist = []

seqrec = sr.seq
for i in seqrec:
    slist.append(i)

qualrec = qr.letter_annotations["phred_quality"]
qtotal = 0
for i in qualrec:
    qlist.append(i)
    qtotal += i
avgqual = float(qtotal)/len(qr)

print "Average quality of " + contigname + " = ", avgqual
#for i, j in zip(slist, qlist):
#    if j < 40:
#        print i, j

belowmins = 0

print "Bases below " + str(MINQUAL) + ":" 
for i, j in enumerate(qlist):
    if j < MINQUAL:
        print i, slist[i], j
        belowmins += 1

seqfile.close()
qualfile.close()

fig = plt.figure(1, figsize=(13,6))
contig_size = len(sr)

gx = [1, contig_size]
gy = [1, 1]

xcoords = range(0, contig_size)

##GC values
gc_values = []
gcx = []

i = 0

while i <= len(sr.seq)+1000:
    gc = GC(sr.seq[i:i+1000])
    gc_values.append(gc)
    gcx.append(i)
    i += 1000

##Plots
plt.title('Contig quality and %GC of ' + contigname)
ax1 = fig.add_subplot(111)
ax1.plot(gx, gy, color='#AF7817', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)

ax1.plot(gcx, gc_values, color='red', linestyle='-', alpha=0.4)
ax1.set_xlabel('Length of contig (bp)')
ax1.set_ylabel('G+C%')

ax2 = ax1.twinx()
ax2.plot(xcoords, qlist, color='#9933CC', linestyle='-', alpha=0.4)
ax2.fill_between(xcoords, qlist, facecolor='#660066', alpha=0.1)
ax2.axis([0, contig_size, 0, 100])
ax2.set_ylabel('Phred quality values')
plt.annotate('Average quality =' + str(fmt(avgqual)), xy=(0.3, 0.4), xycoords='axes fraction')
plt.annotate('Number of bases below phred quality ' + str(MINQUAL) + ' = ' + str(belowmins), xy=(0.3, 0.3), xycoords='axes fraction')

plt.show()

